#!/usr/bin/env python
# -*- coding: utf-8 -*-


from fg_ccmp_common import *
from datasets.utils.data import *
from datasets.utils.plot import *
from datasets.utils.fc_managers import *
from datasets.utils.fg_dataset_mold_base import FG_DatasetMoldBase
from metlib.datetime import parse_datetime
from datakeeper.dataset_mold import GetFileError, GetDataError


class FG_CCMP_DatasetBase(FG_DatasetMoldBase):
    uri_fields = ['dataset', 'subset', 'varname', 'time', 'lon_lat']
    zipped_fields = set()
    uri_parser = parse_ccmp_uri
    allow_RY_download_json = True
    zh_subset_d = ccmp_zh_subset_d
    zh_varname_d = ccmp_zh_varname_d

    def __init__(self, name, uri, info, *args, **kwargs):
        super(FG_CCMP_DatasetBase, self).__init__(name, uri, info, *args, **kwargs)
        self.begdt = parse_datetime(info.get('begdt'))
        self.enddt = parse_datetime(info.get('enddt'))
        self.years = info.get('years')
        self.zh_name = info.get('zh_name', name)
        self.root_path = info.get('root_path', '')
        self.desc = info.get('desc', '')

    def uri_is_standard(self, uri):
        urid = self.parse_uri(uri)
        zipped_fields = set([k for k, v in urid.iteritems() if v == '*'])
        return zipped_fields == self.zipped_fields


class FG_CCMP_PT_Dataset(FG_CCMP_DatasetBase):
    def __init__(self, name, uri, info, *args, **kwargs):
        super(FG_CCMP_PT_Dataset, self).__init__(name, uri, info, *args, **kwargs)
        self.varnames = ccmp_pt_varnames
        self.uri_field_values = {
            'varname': self.varnames
        }

    @property
    def schema(self):
        res = {
            "arch": "subset",
            "name": "PT",
            "type": "SV",
            "subs": []
        }

        for vn in self.varnames:
            res['subs'].append({
                "arch": "variable",
                "name": vn,
                "vartype": "T",
                "datetype": "timeseries",
                "formats": ["json", "csv", "hc_figure"],
                "coords": [{
                    "name": "time",
                    "values": self.years,
                    "default": [self.years[-1]]
                }, {
                    "name": "jy_ix",
                    "values": "<%= jy_ix_values %>",
                    "default": "<%= jy_ix_default %>"
                }]
            })
        return res

    @cache_data
    @perm_data
    def get_data(self, uri, *args, **kwargs):
        datatype = kwargs.get('datatype', 'datapack')
        uri_info = kwargs.get('uri_info', None)
        sample = kwargs.get('sample', False)

        if uri_info is None:
            uri_info = parse_ccmp_uri(uri)

        try:
            try:
                if uri_info['timezone'] != 0:
                    utc_offset = uri_info['timezone']
                else:
                    utc_offset = int(kwargs.get('utc_offset', 8))
                toffset = TD('%sh' % utc_offset)
            except ValueError as e:
                raise BadDataParameter(unicode(e))

            if utc_offset != 0:
                uri_info['begdt'] -= toffset
                uri_info['enddt'] -= toffset

            if sample:
                # subset_uri = uri_info['subset'] + '_sample'
                subset_uri = uri_info['subset']
            else:
                subset_uri = uri_info['subset']

            varname = uri_info['varname']
            units = ccmp_units_d.get(varname, '')

            try:
                s3uri = 'ccmp/%s/%s.h5' % (subset_uri, uri_info['jy_ix'])
                fc = raw_small_fc_manager.get(s3uri, when_not_exist=['fetch'])
                pt_hdf = pd.HDFStore(fc.filepath, mode='r')
                where_cond = 'index >= "%s" & index < "%s" & columns==%s' % (uri_info['begdt'], uri_info['enddt'], varname)
                final_df = pt_hdf.select(key='df', where=where_cond, auto_close=True)
                dts = (final_df.index + toffset).tolist()
                final_df.index = dts
                final_df.index.name = 'datetime'
            except Exception as e:
                # 没有取到数据
                return {}

            if sample:
                # 样本策略
                final_df['month'] = final_df.index.month
                final_df[varname+'_sample'] = pd.rolling_mean(final_df[varname], 24)
                w = final_df.month.isin([2, 3, 5, 6, 8, 9, 11, 12])
                final_df[varname][w] = final_df[varname+'_sample'][w]
                final_df.drop(['month', varname+'_sample'], axis=1, inplace=True)

            final_pack = {
                "type": "datapack",
                "uri": uri,
                "tags": {
                    uri: uri
                },
                "contents": {
                    uri: {
                        "type": "dataunit",
                        "uri": uri,
                        "info": {
                            "varname": varname,
                            "zh_varname": ccmp_zh_varname_d.get(varname, varname),
                            "type": "timeseries",
                            "sample": sample,
                        },
                        "data": {
                            "values": final_df[varname].values,  # FIXME
                            "units": units,
                            "dts": dts,
                            "suggest_range": ccmp_suggest_range_d.get(varname, (None, None)),
                            "begdt": dts[0],
                            "interval": 6 * 3600,
                            "utc_offset": utc_offset,
                        }
                    }
                }
            }

            if datatype == 'datapack':
                res = final_pack
            elif datatype == 'dataframe':
                res = final_df
            else:
                res = final_pack['contents'][uri]
        except Exception as e:
            raise GetDataError(unicode(e))

        return res

    @support_csv
    @support_packunitjson
    def get_file(self, uri, dest, request=None, *args, **kwargs):
        pass


class FG_CCMP_PY_Dataset(FG_CCMP_DatasetBase):
    def __init__(self, name, uri, info, *args, **kwargs):
        super(FG_CCMP_PY_Dataset, self).__init__(name, uri, info, *args, **kwargs)
        self.varnames = ccmp_py_varnames
        self.uri_field_values = {
            'varname': self.varnames
        }

    @property
    def schema(self):
        res = {
            "arch": "subset",
            "name": "PY",
            "type": "SV",
            "subs": []
        }
        for vn in self.varnames:
            if vn in ['mean', 'dist']:
                time_can_poly = True
            else:
                time_can_poly = False

            if vn == 'mean':
                formats = ['json']
            elif vn == 'dist':
                formats = ['json', 'hc_figure']
            elif vn == 'rose':
                formats = ['json', 'wws', 'hc_figure']
            else:
                formats = ['json']

            res['subs'].append({
                "arch": "variable",
                "name": vn,
                "type": "V",
                "vartype": "T",
                "datatype": "yearlystat",
                "formats": formats,
                "zipped": False,
                "coords": [{
                    "name": "time",
                    "values": self.years,
                    "default": [self.years[-1]],
                    "can_poly": time_can_poly,
                }, {
                    "name": "jy_ix",
                    "values": "<%= jy_ix_values %>",
                    "default": "<%= jy_ix_default %>"
                }]
            })
        return res


    @cache_data
    @pack_data
    @perm_data
    def get_data(self, uri, *args, **kwargs):
        datatype = kwargs.get('datatype', 'datapack')
        uri_info = kwargs.get('uri_info', None)
        sample = kwargs.get('sample', False)

        if uri_info is None:
            uri_info = parse_ccmp_uri(uri)

        try:
            if sample:
                subset_uri = uri_info['subset']
            else:
                # subset_uri = uri_info['subset'] + '_sample'
                subset_uri = uri_info['subset']

            raw_fc_manager = raw_tiny_fc_manager
            if uri_info['varname'] == 'mean':
                s3uri = 'ccmp/%s/%s/%s.nc' % (subset_uri, uri_info['varname'], uri_info['time'])
                raw_fc_manager = raw_medium_fc_manager
            elif uri_info['varname'] in ['rose', 'dist']:
                s3uri = 'ccmp/%s/%s/%s/%s.json' % (subset_uri,
                                                   'rose',
                                                   uri_info['time'],
                                                   uri_info['jy_ix'])
            else:
                pass

            fc = raw_fc_manager.get(s3uri, when_not_exist=['fetch'])
            fname = fc.filepath
            if uri_info['varname'] == 'mean':  # 特殊处理
                ccmp_py = CCMP_PY_Handler(fname)
                grid_lat, grid_lon = [float(i) for i in uri_info['jy_ix'].split('_')]
                raw_pack = ccmp_py.get_py_data(grid_lon, grid_lat)
                ccmp_py.close()
            elif uri_info['varname'] in ['rose', 'dist']:
                with open(fname) as fp:
                    raw_pack = json.load(fp, strict=False)

                # add weibull_a, weibull_k
                weibull_a = raw_pack['contents']['data']['data']['a']
                weibull_k = raw_pack['contents']['data']['data']['k']
                raw_pack['contents']['data']['data']['weibull_a'] = weibull_a
                raw_pack['contents']['data']['data']['weibull_k'] = weibull_k
                if uri_info['varname'] == 'dist':
                    # dist_uri = 'ccmp/%s/%s/%s/%s' % (uri_info['subset_uri'],
                    #                                  uri_info['varname'],
                    #                                  uri_info['time'],
                    #                                  uri_info['jy_ix'])
                    # # raw_pack['uri'] = dist_uri
                    # raw_pack['contents']['data']['uri'] = dist_uri
                    raw_pack['contents']['data']['info']['varname'] = uri_info['varname']
                    raw_pack['contents']['data']['info']['type'] = 'winddist'

            thetag = None
            for tag, unituri in raw_pack['tags'].iteritems():
                if uri_info['varname'] == 'dist':
                    unituri = unituri.replace('rose', uri_info['varname'], 1)
                if unituri == uri:
                    thetag = tag
                    break
            if thetag:
                res = raw_pack['contents'][thetag]
                varname = uri_info['varname']
                if varname == 'rose' and datatype == 'rawrose':
                    res['data']['units'] = ccmp_units_d.get(varname, u'‰')
                else:
                    if varname == 'rose':
                        data = res['data']
                        data['wsbins'] = ['<3', '3-5', '5-7', '7-9', '9-11', '11-13', '>13']
                        data['wsbinnum'] = len(data['wsbins'])
                        v = data['values']
                        new_values = np.zeros((data['wsbinnum'], data['wdbinnum']))
                        new_values[0] = np.sum(v[0:3], axis=0)
                        new_values[1] = np.sum(v[3:5], axis=0)
                        new_values[2] = np.sum(v[5:7], axis=0)
                        new_values[3] = np.sum(v[7:9], axis=0)
                        new_values[4] = np.sum(v[9:11], axis=0)
                        new_values[5] = np.sum(v[11:13], axis=0)
                        new_values[6] = np.sum(v[13:], axis=0)
                        data['values'] = new_values.tolist()
                        if varname == 'rose':
                            data['units'] = ccmp_units_d.get(varname, u'‰')
                    if varname == 'dist':
                        res['data']['units'] = '%'
                # res['contents']['data']['data'] = data
                try:
                    res['info'].update(raw_pack['common']['info'])
                except Exception:
                    pass
                amend_dataunit(res, info={
                    "info": {
                        "sample": sample,
                        "zh_varname": ccmp_zh_varname_d.get(varname, varname),
                        "level": '10'
                    }
                })
                res["uri"] = uri
            else:
                raise GetDataError('Cannot get data: %s with kwargs %s' % (uri, kwargs))
        except GetDataError as e:
            raise e
        except Exception as e:
            raise GetDataError(unicode(e))

        return res

    @support_wws
    @support_packunitjson
    def get_file(self, uri, dest, request=None, *args, **kwargs):
        pass


class FG_CCMP_PM_Dataset(FG_CCMP_DatasetBase):
    zipped_fields = {'varname'}

    def __init__(self, name, uri, info, *args, **kwargs):
        super(FG_CCMP_PM_Dataset, self).__init__(name, uri, info, *args, **kwargs)
        self.varnames = ccmp_pm_varnames
        self.uri_field_values = {
            'varname': self.varnames
        }

    @property
    def schema(self):
        res = {
            "arch": "subset",
            "name": "PM",
            "type": "SV",
            "subs": []
        }

        for vn in self.varnames:
            res['subs'].append({
                "arch": "variable",
                "name": vn,
                "vartype": "T",
                "datatype": "monthlystat",
                "formats": ["json", "hc_figure"],
                "zipped": True
            })
        return res

    @cache_data
    @pack_data
    @perm_data
    def get_data(self, uri, *args, **kwargs):
        datatype = kwargs.get('datatype', 'datapack')
        uri_info = kwargs.get('uri_info', None)
        sample = kwargs.get('sample', False)

        if uri_info is None:
            uri_info = parse_ccmp_uri(uri)

        try:
            if sample:
                subset_uri = uri_info['subset']
            else:
                # subset_uri = uri_info['subset'] + '_sample'
                subset_uri = uri_info['subset']

            s3uri = 'ccmp/%s/%s/%s.nc' % (subset_uri, 'mean', uri_info['time'])
            fc = raw_large_fc_manager.get(s3uri, when_not_exist=['fetch'])
            fname = fc.filepath
            pm = CCMP_PM_Handler(fname)
            grid_lat, grid_lon = [float(i) for i in uri_info['jy_ix'].split('_')]
            raw_pack = pm.get_pm_data(grid_lon, grid_lat, uri_info['varname'])
            pm.close()
            thetag = None
            for tag, unituri in raw_pack['tags'].iteritems():
                if unituri == uri:
                    thetag = tag
                    break
            if thetag:
                res = raw_pack['contents'][thetag]
                amend_dataunit(res, info={
                    "info": {
                        "sample": sample
                    }
                })
            else:
                raise GetDataError('Cannot get data: %s with kwargs %s' % (uri, kwargs))
        except GetDataError as e:
            raise e
        except Exception as e:
            raise GetDataError(unicode(e))

        return res

    @support_packunitjson
    def get_file(self, uri, dest, request=None, *args, **kwargs):
        pass


class FG_CCMP_RY_Dataset(FG_CCMP_DatasetBase):
    def __init__(self, name, uri, info, *args, **kwargs):
        super(FG_CCMP_RY_Dataset, self).__init__(name, uri, info, *args, **kwargs)
        self.varnames = ccmp_ry_varnames
        self.uri_field_values = {
            'varname': self.varnames
        }

    @property
    def schema(self):
        res = {
            'arch': 'subset',
            'name': 'RY',
            'type': 'SV',
            'subs': []
        }
        for vn in self.varnames:
            res['subs'].append({
                'arch': 'variable',
                'name': vn,
                'vartype': 'T',
                'datatype': 'yearlystat',
                'formats': ['json', 'rect_figure', 'kmz'],
                'zipped': False,
                "coords": [{
                    "name": "time",
                    "values": self.years,
                    "default": [self.years[-1]]
                }, {
                    "name": "jy_ix",
                    "values": "<%= jy_ix_values %>",
                    "default": "<%= jy_ix_default %>",
                }]
            })
        return res

    @cache_data(valid_params={'dataunit': ['datatype', 'user']})
    @provide_figinfo
    @perm_data
    def get_data(self, uri, *args, **kwargs):
        datatype = kwargs.get('datatype', 'figinfo')
        uri_info = kwargs.get('uri_info', None)
        sample = kwargs.get('sample', False)

        if uri_info is None:
            uri_info = parse_ccmp_uri(uri)

        try:
            varname = uri_info['varname']
            jy_ix = uri_info['jy_ix']

            if sample:
                # subset_uri = uri_info['subset'] + '_sample'
                subset_uri = uri_info['subset']
            else:
                subset_uri = uri_info['subset']

            s3uri = 'ccmp/PY/mean/%s.nc' % uri_info['time']
            fc = raw_medium_fc_manager.get(s3uri, when_not_exist=['fetch'])
            fname = fc.filepath

            jys, ixs = jy_ix.split('_')
            lat1, lat2 = [float(j) for j in jys.split(':')]
            lon1, lon2 = [float(i) for i in ixs.split(':')]

            # 区域扩充
            lat_delta = self.info.get('lat_delta')
            lon_delta = self.info.get('lon_delta')
            lat1 -= 1.0 * lat_delta
            lon1 -= 1.0 * lon_delta
            lat2 += 1.5 * lon_delta
            lon2 += 1.5 * lon_delta

            # jy1, jy2, ix1, ix2
            ccmp_rect = get_ccmp_rect(lon1, lat1, lon2, lat2, self.info)
            ix1, jy1 = lonlat_to_ccmp_ij(ccmp_rect.get('lon1'), ccmp_rect.get('lat1'))
            ix2, jy2 = lonlat_to_ccmp_ij(ccmp_rect.get('lon2'), ccmp_rect.get('lat2'))
            ix2 += 1
            jy2 += 1

            ncf = Dataset(fname, 'r')
            v = ncf.variables[uri_info['varname']]
            data = v[0, jy1:jy2, ix1:ix2]
            lat = ncf.variables['lat'][jy1:jy2]
            lon = ncf.variables['lon'][ix1:ix2]
            ncf.close()

            res = {
                "uri": uri,
                "type": "dataunit",
                "info": {
                    "varname": varname,
                    "type": "rect_yearlystat",
                    "sample": sample,
                },
                "coords": {
                    "dataset": uri_info['dataset'],
                    "subset": uri_info['subset'],
                    "varname": uri_info['varname'],
                    "time": uri_info['time'],
                    "jy_ix": uri_info['jy_ix']
                },
                "data": {
                    "lat": lat,
                    "lon": lon,
                    "values": data,
                    "suggest_range": ccmp_suggest_range_d.get(varname, (None, None)),
                    "units": ccmp_units_d.get(varname, ''),
                }
            }

        except GetDataError as e:
            raise e
        except Exception as e:
            raise GetDataError(unicode(e))

        return res

    def get_file(self, uri, dest, request=None, *args, **kwargs):
        return self.RY_get_file(uri, dest, request=request, *args, **kwargs)


class FG_CCMP_Dataset(FG_CCMP_DatasetBase):
    mold_name = 'CCMP'
    subset_classes = [
        {'name': 'PointYearlyStat', 'uri': 'PY', 'class': FG_CCMP_PY_Dataset},
        {'name': 'PointMonthlyStat', 'uri':'PM', 'class': FG_CCMP_PM_Dataset},
        {'name': 'PointTimeseries', 'uri': 'PT', 'class': FG_CCMP_PT_Dataset},
        # {'name': 'PointDiurnalStat', 'uri':'PD', 'class': FG_Merra_PD_Dataset},
        {'name': 'RectYearlyStat', 'uri':'RY', 'class': FG_CCMP_RY_Dataset}
    ]
    tags = ['wind', 'ccmp', 'meteorology']

    def __init__(self, name, uri, info, *args, **kwargs):
        super(FG_CCMP_Dataset, self).__init__(name, uri, info, *args, **kwargs)

    @property
    def schema(self):
        res = {
            "arch": "dataset",
            "mold": self.mold_name,
            "type": "DSV",
            "name": self.name,
            "zh_name": self.zh_name,
            "desc": self.desc,
            "common": {
                "coords": [{
                    "name": "time",
                    "values": self.years,
                    "default": [self.years[-1]],
                    "can_poly": True,
                }, {
                    "name": "jy_ix",
                    "values": "<%= jy_ix_values %>",
                    "default": "<%= jy_ix_default %>",
                }]
            },
            "subset_dict": ccmp_zh_subset_d,
            "varname_dict": ccmp_zh_varname_d,
            "subs": []
        }
        for class_info in self.subset_classes:
            sub_schema = self.uri_subsets[class_info['uri']].schema
            res['subs'].append(sub_schema)
        return res

    def lookup(self, info, *args, **kwargs):
        try:
            results = []
            if info['type'] == 'point':
                lon = float(info.get('lon'))
                lat = float(info.get('lat'))
                ccmp_latlon = latlon_to_ccmp_latlon(lat, lon, self.info)
                if ccmp_latlon is not None:
                    results.append({
                        "dataset": self.name,
                        "subsets": ["PT", "PY", 'PM'],
                        "jy_ix_values": [ccmp_latlon],
                        "jy_ix_default": [ccmp_latlon]
                    })
            elif info['type'] == 'rect':
                lon1 = float(info.get('lon1'))
                lat1 = float(info.get('lat1'))
                lon2 = float(info.get('lon2'))
                lat2 = float(info.get('lat2'))
                rect_info = get_ccmp_rect(lon1, lat1, lon2, lat2, self.info)
                if rect_info is not None:
                    results.append({
                        "dataset": self.name,
                        "subsets": ["RY"],
                        "jy_ix_values": [rect_info['jy_ix']],
                        "jy_ix_default": [rect_info['jy_ix']]
                    })
            else:
                pass
        except Exception as e:
            raise LookupError(u'[%s] %s' % (type(e), unicode(e)))
        return results

    def get_data(self, uri, *args, **kwargs):
        try:
            if 'uri_info' in kwargs:
                uri_info = kwargs['uri_info']
            else:
                uri_info = parse_ccmp_uri(uri)
            subset = self.uri_subsets.get(uri_info['subset'])
            if subset is None:
                return
            res = subset.get_data(uri, *args, **kwargs)
            return res
        except GetDataError as e:
            raise e
        except Exception as e:
            raise GetDataError(u'[%s] %s' % (type(e), unicode(e)))

    @protect_hack
    def get_file(self, uri, dest, request=None, *args, **kwargs):
        try:
            # 防止用户插入奇怪的参数, 如果数据集内部调用时有敏感参数, 则需在此做防备.
            uri_info = parse_ccmp_uri(uri)
            subset = self.uri_subsets.get(uri_info['subset'])
            if subset is None:
                raise ValueError("No such subset in ccmp: %s" % uri_info['subset'])

            getres = subset.get_file(uri, dest, request=request, uri_info=uri_info, *args, **kwargs)
            return getres
        except GetFileError as e:
            raise e
        except Exception as e:
            raise GetFileError(u'[%s] %s' % (type(e), unicode(e)))

    def is_free(self, uri=None, *args, **kwargs):
        return False
